import { PluginObj, transformFromAstSync } from '@babel/core';
import parser from '@babel/parser';
import template from '@babel/template';
import { isObjectExpression } from '@babel/types';
import prettier from 'prettier';

const sourceCode = `
import { Module } from '@nestjs/common';

@Module({})
export class AaaModule {}
`;

function myPlugin(): PluginObj {
    return {
        visitor: {
            Program(path: any) {
                // 首先，从上到下找 ImportDeclaration，直到最后一个，然后在最后面插入一个 import 语句。
                let index = 0;
                while (path.node.body[index].type === 'ImportDeclaration') {
                    index++;
                }
                template.statement
                // 这里用 @babel/template 包的 api 来创建这个 ast。
                const ast = template.statement("import { AaaController } from './aaa.controller';")()
                // 然后就是在 controllers 数组插入一个元素。
                path.node.body.splice(index, 0, ast);
            },
            Decorator(path: any) {
                const decoratorName = path.node.expression.callee.name
                if (decoratorName !== 'Module') {
                    return;
                }
                const obj = path.node.expression.arguments[0];
                const controllers = obj.properties.find((item: any) => item.key.name === 'controllers');
                if (!controllers) {
                    const expression = template.expression('{controllers: [AaaController]}')();
                    if (isObjectExpression(expression)) {
                        obj.properties.push(expression.properties[0]);
                    }
                } else {
                    const property = template.expression('AaaController')();
                    controllers.value.elements.push(property);
                }
            }
        }
    }
}

const ast = parser.parse(sourceCode, {
    // 设置 sourceType 为 module 就是按照 es module 来解析
    sourceType: 'module',
    // 这里要指定 decoratos 的装饰器 语法插件，不然解析不了装饰器。
    plugins: ["decorators"]
})

const res = transformFromAstSync(ast, sourceCode, {
    plugins: [myPlugin],
    retainLines: true // 它会保留原本的行列号。
});


(async function () {
    const formatedCode = await prettier.format(res?.code!, {
        filepath: 'aaa.ts'
    })
    console.log(formatedCode);
})()

